# Testing interaction unit.
# 23-11-01

import os
import sys
import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk
import mindspore as ms
import mindspore.nn as ms_nn

from jax import Array, jit
from typing import Optional, Union, List
from jax.nn.initializers import lecun_normal, normal, zeros, ones
from mindspore import load_param_into_net, Tensor, nn, context

sys.path.append(os.path.dirname(sys.path[0]))

# setting context & hyper parameters
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")

from cybertron.model.interaction.molct_interaction import InteractionUnit
from ms_cybertron.interaction import NIUWithEdge as ms_InteractionUnit

EPSILON = 1e-3
rng = jax.random.PRNGKey(42)

# generate data for test
np.random.seed(42)
node_vec = np.random.uniform(size=(3, 5, 4)) # (B, A, F)
node_mask = np.sum(node_vec, axis=-1) < 2 # (B, A)

edge_vec = np.random.uniform(size=(3, 5, 5, 4)) # (B, A, A, F)
eye = np.eye(5).reshape(1, 5, 5, 1)
eye_mask = np.array(eye, dtype=np.bool_)
edge_vec = np.where(eye_mask, eye, edge_vec)

edge_mask = np.sum(edge_vec, axis=-1) < 2 # (B, A, A)
eye = np.eye(5).reshape(1, 5, 5)
eye_mask = np.array(eye, dtype=np.bool_)
edge_mask = np.where(eye_mask, eye_mask, edge_mask)

edge_cutoff = np.random.uniform(size=(3, 5, 5)) # (B, A, A)
edge_cutoff = np.where(eye_mask, eye, edge_cutoff)

is_edge_update = False

print("=================Test NeuralInteractionUnit====================")

def niu_fn(a, b, c, d, e, training=False):
    niu = InteractionUnit(dim_feature=4, 
                          n_heads=2, 
                          dim_outer_pdct=8, 
                          activation='silu', 
                          is_edge_update=is_edge_update)(a, b, c, d, e) # type: ignore
    return niu

niu_fn = jax.vmap(niu_fn, in_axes=(0, 0, 0, 0, 0))
ts_niu_fn = hk.transform(niu_fn, apply_rng=True)
niu_params = ts_niu_fn.init(rng, node_vec, node_mask, edge_vec, edge_mask, edge_cutoff)
node_out, edge_out = jit(ts_niu_fn.apply)(niu_params, rng, node_vec, node_mask, edge_vec, edge_mask, edge_cutoff)
# print(f"niu_node_out: shape of {node_out.shape}, value of \n{node_out}")
print(f"niu_edge_out: shape of {edge_out.shape}, value of \n{edge_out}")

_name = 'neural_interaction_unit/~/'
_params = hk.data_structures.to_mutable_dict(niu_params)
ms_niu_fn = ms_InteractionUnit(dim_feature=4,
                               n_heads=2, 
                               dim_outer_product=8, 
                               activation='silu',
                               edge_update=is_edge_update,
                               max_cycles=1,)
ms_params = {}
ms_params['positional_embedding.norm.gamma'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/norm']['scale']), ms.float32))
ms_params['positional_embedding.norm.beta'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/norm']['offset']), ms.float32))
ms_params['positional_embedding.x2q.weight'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/q_gen']['w']).T, ms.float32))
ms_params['positional_embedding.x2k.weight'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/k_gen']['w']).T, ms.float32))
ms_params['positional_embedding.x2v.weight'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/v_gen']['w']).T, ms.float32))
ms_params['multi_head_attention.output.weight'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}multihead_attention/linear_output']['w']).T, ms.float32))
ms_params['norm.beta'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}node_norm']['offset']), ms.float32))
ms_params['norm.gamma'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}node_norm']['scale']), ms.float32))
if is_edge_update:
    ms_params['edge_norm.beta'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}edge_norm']['offset']), ms.float32))
    ms_params['edge_norm.gamma'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}edge_norm']['scale']), ms.float32))
    
    _name = 'neural_interaction_unit/~/hyperformer_pair_block/'
    ms_params['hyperformer_pair_block.outer_product.left_projection_weights'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}outer_product']['left_weights']).T, ms.float32))
    ms_params['hyperformer_pair_block.outer_product.right_projection_weights'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}outer_product']['right_weights']).T, ms.float32))
    ms_params['hyperformer_pair_block.outer_product.left_projection_biases'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}outer_product']['left_bias']), ms.float32))
    ms_params['hyperformer_pair_block.outer_product.right_projection_biases'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}outer_product']['right_bias']), ms.float32))
    ms_params['hyperformer_pair_block.outer_product.linear_output_weights'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}outer_product']['output_weights']).T, ms.float32))
    ms_params['hyperformer_pair_block.outer_product.o_biases'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}outer_product']['output_bias']), ms.float32))
    ms_params['hyperformer_pair_block.outer_product.layer_norm.gamma'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}outer_product/~/norm_fn']['scale']), ms.float32))
    ms_params['hyperformer_pair_block.outer_product.layer_norm.beta'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}outer_product/~/norm_fn']['offset']), ms.float32))
    ms_params['hyperformer_pair_block.pair_transition.transition1_weights'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}transition/~/transition_1']['w']).T, ms.float32))
    ms_params['hyperformer_pair_block.pair_transition.transition1_biases'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}transition/~/transition_1']['b']), ms.float32))
    ms_params['hyperformer_pair_block.pair_transition.transition2_weights'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}transition/~/transition_2']['w']).T, ms.float32))
    ms_params['hyperformer_pair_block.pair_transition.transition2_biases'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}transition/~/transition_2']['b']), ms.float32))
    ms_params['hyperformer_pair_block.pair_transition.layer_norm.gamma'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}transition/~/norm']['scale']), ms.float32))
    ms_params['hyperformer_pair_block.pair_transition.layer_norm.beta'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}transition/~/norm']['offset']), ms.float32))
load_param_into_net(ms_niu_fn, ms_params)

ms_node_vec = Tensor(node_vec, ms.float32)
ms_node_mask = Tensor(node_mask, ms.bool_)
ms_edge_vec = Tensor(edge_vec, ms.float32)
ms_edge_mask = Tensor(edge_mask, ms.bool_)
ms_edge_cutoff = Tensor(edge_cutoff, ms.float32)

ms_node_out, ms_edge_out = ms_niu_fn(node_vec=ms_node_vec,
                                     node_emb=ms_node_vec,
                                     node_mask=ms_node_mask,
                                     edge_vec=ms_edge_vec,
                                     edge_emb=ms_edge_vec,
                                     edge_mask=ms_edge_mask,
                                     edge_cutoff=ms_edge_cutoff,) # type: ignore

# print(f"ms_niu_node_out: shape of {ms_node_out.shape}, value of \n{ms_node_out}")
print(f"ms_niu_edge_out: shape of {ms_edge_out.shape}, value of \n{ms_edge_out}")

niu_node_same = np.allclose(node_out, ms_node_out.asnumpy(), atol=EPSILON)
print(f"niu_node is the same: {niu_node_same} in accuracy of {EPSILON}")
niu_node_diff = node_out - ms_node_out.asnumpy()
niu_node_norm = np.linalg.norm(niu_node_diff)
print(f"norm of diff: {niu_node_norm}")

niu_edge_same = np.allclose(edge_out, ms_edge_out.asnumpy(), atol=EPSILON)
print(f"niu_edge is the same: {niu_edge_same} in accuracy of {EPSILON}")
niu_edge_diff = edge_out - ms_edge_out.asnumpy()
niu_edge_norm = np.linalg.norm(niu_edge_diff)
print(f"norm of diff: {niu_edge_norm}")